import jwt
import time
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.backends import default_backend


def generate_key_pair():
    """生成RSA密钥对并保存到文件（只需运行一次）"""
    from cryptography.hazmat.primitives.asymmetric import rsa

    # 生成私钥
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
        backend=default_backend()
    )

    # 保存私钥
    with open("private_key.pem", "wb") as f:
        f.write(private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        ))

    # 生成并保存公钥
    public_key = private_key.public_key()
    with open("public_key.pem", "wb") as f:
        f.write(public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        ))
    print("密钥对生成成功：private_key.pem 和 public_key.pem")


def load_private_key():
    """加载私钥（用于生成Token）"""
    with open("private_key.pem", "rb") as f:
        private_key = serialization.load_pem_private_key(
            f.read(),
            password=None,
            backend=default_backend()
        )
    return private_key


def load_public_key():
    """加载公钥（用于验证Token）"""
    with open("public_key.pem", "rb") as f:
        public_key = serialization.load_pem_public_key(
            f.read(),
            backend=default_backend()
        )
    return public_key


def generate_token(user_id, private_key):
    """用私钥生成JWT"""
    payload = {
        "sub": str(user_id),  # 标准字段：用户标识
        "exp": int(time.time()) + 3600 * 24 * 7,  # 7天有效期
        "iat": int(time.time())
    }
    return jwt.encode(
        payload,
        private_key,
        algorithm="RS256"
    )


def verify_token(token, public_key):
    """用公钥验证JWT"""
    try:
        payload = jwt.decode(
            token,
            public_key,
            algorithms=["RS256"],
            options={"verify_exp": True}
        )
        return {"valid": True, "user_id": payload["sub"]}
    except jwt.ExpiredSignatureError:
        return {'valid': False, 'msg': 'Token已过期'}
    except jwt.InvalidTokenError as e:
        return {'valid': False, 'msg': str(e)}

private_key = load_private_key()
public_key = load_public_key()

if __name__ == '__main__':
    # 第一步：生成密钥对（首次运行时取消注释，生成后注释掉）
    # generate_key_pair()

    # 第二步：加载密钥


    # 第三步：生成并验证Token
    token = generate_token(1, private_key)
    print("生成的Token：", token)

    result, msg = verify_token(token, public_key)
    print("验证结果：", (result, msg))
